import os
from mmdeploy.apis import torch2onnx
from mmdeploy.backend.sdk.export_info import export2SDK

img = '/home/lixiang/data/balloon/train/9330497995_4cf0438cb6_k.jpg'
work_dir = '/home/lixiang/.config/labelapp/output/YYWahrFLTiiQskmIS9REg'
model_cfg = os.path.join( work_dir, 'faster-rcnn_r50-caffe_fpn_ms-2x_coco.py' )
model_checkpoint = os.path.join( work_dir, 'epoch_24.pth' )
save_file = 'end2end.onnx'
deploy_cfg = '/home/lixiang/projects/openmmlab/mmdeploy/configs/mmdet/detection/detection_onnxruntime_dynamic.py'
device = 'cpu'

# 1. convert model to onnx
torch2onnx(img, work_dir, save_file, deploy_cfg, model_cfg, model_checkpoint, device)